Skip to content

feat(domain): benchmark all differentiable solvers in the loop#116

Draft
andrinr wants to merge 25 commits into
fix/ns-grid-recurrent-statefrom
feat/ns-grid-solver-in-loop
Draft

feat(domain): benchmark all differentiable solvers in the loop#116
andrinr wants to merge 25 commits into
fix/ns-grid-recurrent-statefrom
feat/ns-grid-solver-in-loop

Conversation

@andrinr

@andrinr andrinr commented Jul 24, 2026

Copy link
Copy Markdown
Collaborator

What this PR establishes

This PR adds one common 2D solver-in-the-loop neural-correction benchmark for
all six differentiable Navier–Stokes grid solvers:

  • JAX-CFD
  • INS.jl
  • PhiFlow
  • PICT
  • Warp-NS
  • XLB

It is stacked on the opt-in canonical recurrent-state contract in #121.
JAX-CFD, INS.jl, PhiFlow, and XLB carry their native checkpoint through
training, evaluation, plotting, finite-difference validation, and
differentiation. PICT and Warp-NS are velocity-complete. Both canonical
velocity and native checkpoint state are stopped in the paired stop-gradient
control.

The benchmark makes two deliberately separate measurements:

  1. absolute correctability, measured against a declared reference;
  2. temporal-credit value, measured within each solver by pairing full
    differentiation through the autoregressive solver–corrector recurrence with
    the identical stopped-state recurrence.

The second comparison is the primary differentiability result. The benchmark
is therefore best read as a recurrent credit-assignment audit: does
backpropagating through several corrected solver transitions improve a policy
over local one-step supervision? It is not a claim that a VJP alone makes the
corrected solver accurate, nor that one numerical reference is universally
neutral.

Reference design

The nonlinear task is repeated against two independently discretized targets:

Reference Production target Convergence audit
dealiased pseudo-spectral 128², dt/4 256², dt/8
conservative finite volume MUSCL/minmod + Rusanov + five-point diffusion, 128², dt/4 256², dt/8

Both audits start from the same Fourier-prolonged continuum initial condition
and compare seeds 0 and 100 at frames 1, 8, 24, and 36. Training is rejected if
the p95 128²-versus-256² discrepancy exceeds 0.5%.

Audit Median / p95 / maximum relative discrepancy Gate
pseudo-spectral 128² vs 256² p95 9.89e-8; max 1.01e-7 pass
finite-volume 128² vs 256² median 0.121%; p95 0.160%; max 0.168% pass
pseudo-spectral vs finite-volume, all held-out frames median 0.162%; p95 0.207%; max 0.224% diagnostic
pseudo-spectral vs finite-volume, final frame median 0.196%; p95 0.222%; max 0.224% diagnostic

The finite-volume method uses an FFT only to solve the periodic discrete
Poisson equation. Transport and diffusion use conservative grid stencils, so
this is not a second spectral time integrator.

Reference convergence and conclusion sensitivity

Reference vorticity fields and differences

Reference disagreement across all held-out ICs

Common paired protocol

Every cell uses the same 32² candidate grid, viscosity, physical time step,
16 training ICs, eight held-out ICs, three paired model seeds, 200 optimizer
updates, recurrent training to t=1.92, and evaluation to t=2.88.

The learned component is the same zero-initialized Equinox periodic residual
CNN (56,098 parameters). Only the differentiable solver transition and its VJP
change between cells.

Training uses truncated backpropagation through eight consecutive correction
intervals. Each interval advances four native solver steps (4 × 0.02 = 0.08),
applies the corrector, and feeds both the corrected velocity and native solver
checkpoint into the next interval. There is no teacher forcing inside the
sampled window. The loss averages all eight corrected states and adds a
0.1-weight terminal loss on the last provisional solver state.

Held-out evaluation is fully free-running for 36 solver–corrector intervals:
after the initial condition, no reference field is fed back into the rollout.

The paired modes are:

  1. full temporal credit: gradients flow through eight solver–corrector
    transitions, including solver velocity and native state;
  2. local control: the identical forward recurrence and supervision, but
    recurrent velocity and native state are stopped at every solver output.

The solver-VJP lift is the geometric paired
error_local-control / error_full-temporal-credit ratio. A correction gain
above one separately means that the full-credit corrector beats the
uncorrected solver.

Nonlinear result and reference sensitivity

Absolute error and correction gain remain reference-dependent:

Solver Native final error (spectral / FV) Corrected final error (spectral / FV) Correction gain (spectral / FV)
JAX-CFD 0.06842 / 0.06755 0.1526 / 0.1321 0.885× / 0.967×
INS.jl 0.05117 / 0.05108 0.1492 / 0.1526 0.737× / 0.701×
PhiFlow 0.04461 / 0.04461 0.2921 / 0.1743 0.686× / 0.677×
XLB 0.03587 / 0.03489 0.1117 / 0.1255 0.752× / 0.650×
PICT 0.05016 / 0.05039 0.04448 / 0.03957 1.107× / 1.212×
Warp-NS 0.03568 / 0.03608 0.03809 / 0.03705 0.946× / 0.973×

The native final errors move by at most 2.7% when the reference changes, but
the trained corrected errors move by as much as 40.3% (PhiFlow). PICT is the
only corrector that beats its native solver under both references. This
separates target agreement from non-convex corrector-training sensitivity:
very similar converged targets do not guarantee the same learned endpoint.

The paired solver-VJP conclusion is more robust, but not universal:

Solver Spectral VJP lift [bootstrap 95% CI] FV VJP lift [bootstrap 95% CI] Reference conclusion
JAX-CFD +39.0% [+35.6%, +42.2%] +49.8% [+35.4%, +77.5%] beneficial under both
INS.jl +38.8% [+27.8%, +50.5%] +38.2% [+35.8%, +41.3%] beneficial under both
PhiFlow +37.0% [+21.0%, +48.1%] +45.1% [+38.2%, +54.3%] beneficial under both
XLB +20.5% [+13.1%, +30.2%] +16.2% [+7.2%, +24.6%] beneficial under both
PICT −4.2% [−7.8%, −1.0%] +4.3% [+2.1%, +6.7%] harmful → beneficial
Warp-NS −0.3% [−2.8%, +2.0%] +4.9% [−2.0%, +10.9%] inconclusive under both

JAX-CFD, INS.jl, PhiFlow, and XLB retain a resolved full-VJP benefit under
both references. Warp-NS remains inconclusive. PICT is the important exception:
its paired VJP effect changes from significantly harmful under the spectral
target to significantly beneficial under the finite-volume target, even though
its absolute correction gain remains above one in both tasks.

Pseudo-spectral-reference full fields

Finite-volume-reference full fields

Pseudo-spectral-reference held-out trajectory

Finite-volume-reference held-out trajectory

Recurrent-state admission

All six solvers pass the same nonlinear and Taylor–Green recurrence gate:

Solver Recurrent representation Nonlinear closure p95 Admitted
JAX-CFD native checkpoint 0 yes
INS.jl native checkpoint 2.92e-7 yes
PhiFlow native checkpoint 4.72e-7 yes
PICT velocity-complete 0 yes
Warp-NS velocity-complete 0 yes
XLB native checkpoint 1.75e-3 yes

All-solver recurrence and generalization diagnostics

Analytic and solver-specific controls

Analytic Taylor–Green control

All six solvers pass and train against the analytic Taylor–Green solution.
JAX-CFD, INS.jl, PhiFlow, PICT, and Warp-NS improve absolute error; XLB is
close to neutral. The paired VJP effect is small near this error floor, as
expected.

Taylor–Green full fields

Taylor–Green paired comparison

Taylor–Green held-out trajectory

Solver-specific refined-reference control

Five solvers also complete a method-local task against their own 64²,
half-step target. Those absolute errors are not ranked across solvers because
the targets differ. PICT improves absolute error by 1.119×. The paired
solver-VJP benefit is resolved for JAX-CFD (37.0%), INS.jl (29.8%), XLB
(13.3%), and PICT (4.60%); Warp-NS is inconclusive.

PhiFlow is rejected before training in this control only. Its native state
closes the common-reference task, but the chaotic refined-reference audit has
long-horizon spikes: maximum coarse/fine closure residuals are 0.0101/0.0252
and reach 22.8%/70.7% of the refinement signal.

Self-reference full fields

Self-reference paired comparison

Self-reference held-out trajectory

Interpretation and limits

  • The paired VJP comparison is within solver: it isolates temporal
    differentiation while holding the forward recurrence and supervision fixed.
  • A positive temporal-credit lift says the eight-step training graph helps
    relative to local supervision; it does not by itself say that the learned
    corrector beats the uncorrected solver.
  • Absolute solver and corrected errors remain properties of the chosen target.
    The two-reference sweep tests whether the conclusions survive one independent
    discretization change; it does not establish a universally fair reference.
  • The reference agreement reported here applies to this smooth low-wave-number
    initial-condition distribution, viscosity, resolution, and time horizon.
  • A corrector can benefit from a solver VJP relative to stop-gradient and still
    fail to beat the uncorrected solver. Both quantities are reported.

Implementation state

  • Detect native recurrence only when state is differentiable in both the
    Tesseract input and output schemas.
  • Thread solver state through corrector training, evaluation,
    finite-difference checks, long-closure checks, and self-reference generation.
  • Stop both velocity and native state in the paired control.
  • Add an independent conservative finite-volume reference, exact Fourier
    prolongation for resolution audits, and a gated 128²-versus-256²
    space-time-convergence check.
  • Register the spectral/finite-volume sensitivity study as the canonical
    opt-in optimization/solver_in_loop_reference_sensitivity experiment.
  • Keep the existing drag-optimization registrations unchanged.

Complete merged JSON, plot inputs, full-field/physics plots, GIFs, tables, and
job provenance are in the
artifact index.

Validation and offline execution

All experiments ran offline on Kander through the shared Slurm/Pyxis solver
tooling. No hosted benchmark was triggered; this draft carries
benchmark:none.

  • reference convergence audit: 1697345;
  • all-held-out reference-disagreement audit: 1697383;
  • pseudo-spectral reference production: 1697346, 1697347, 1697348,
    1697349, 1697350, 1697351;
  • finite-volume reference production: 1697352, 1697371, 1697372,
    1697373, 1697374, 1697375;
  • reference-sensitivity merge/render: 1697376;
  • existing nonlinear, Taylor–Green, and method-local controls: see artifact
    index;
  • numerical-source validation at 21db5ff: Ruff and format passed; 522 tests
    passed and three were skipped (1697344);
  • final plot-only layout change: Ruff/format and the offline pre-merge renderer
    passed; hosted Python checks are reported on this PR head.

This remains a draft PR based on #121 until the canonical recurrent-state
contract lands.

@andrinr
andrinr force-pushed the feat/ns-grid-solver-in-loop branch from 2b53888 to 358e231 Compare July 24, 2026 15:02
@andrinr andrinr added benchmark:solver Benchmark only the modified solver and removed benchmark:solver Benchmark only the modified solver labels Jul 24, 2026
@andrinr
andrinr force-pushed the feat/ns-grid-solver-in-loop branch from 7d6a683 to 04b56d8 Compare July 27, 2026 12:51
@andrinr andrinr changed the title feat(domain): add 2D solver-in-the-loop corrector benchmark feat(domain): benchmark all differentiable solvers in the loop Jul 27, 2026
@andrinr
andrinr changed the base branch from main to fix/ns-grid-recurrent-state July 27, 2026 12:51
@andrinr andrinr added the benchmark:none Skip benchmarks (maintainer trusts no answer-changing code) label Jul 27, 2026
@PasteurBot

Copy link
Copy Markdown
Contributor

📊 View the full benchmark results

No benchmarks ran for this PR, so there is no status report.


Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

benchmark:none Skip benchmarks (maintainer trusts no answer-changing code)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants